/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * Licensed under the CANN Open Software License Agreement Version 1.0.
 * See LICENSE in the root of the software repository for full license details.
 */

/**
 * @file gelu_operator.cpp
 * @brief Implementation of GELU operator using Ascend CANN ACL
 */

 #include <iostream>
 #include <vector>
 #include <memory>
 #include <fstream>
 #include <system_error>
 #include <cassert>
 
 #include "acl/acl.h"
 #include "aclnn_gelu.h"
 
 namespace {
 
 // Constants
 constexpr int32_t DEFAULT_DEVICE_ID = 0;
 constexpr aclDataType DATA_TYPE = ACL_FLOAT16;
 constexpr size_t FLOAT16_SIZE = sizeof(aclFloat16);
 
 // Error codes
 enum class GeluError {
     SUCCESS = 0,
     ACL_INIT_FAILURE,
     DEVICE_SET_FAILURE,
     STREAM_CREATE_FAILURE,
     MEMORY_ALLOCATION_FAILURE,
     MEMORY_COPY_FAILURE,
     TENSOR_CREATION_FAILURE,
     WORKSPACE_QUERY_FAILURE,
     OPERATOR_EXECUTION_FAILURE,
     STREAM_SYNC_FAILURE,
     FILE_OPERATION_FAILURE
 };
 
 // Logger class
 class Logger {
 public:
     static void Info(const std::string& message) {
         std::cout << "[INFO] " << message << std::endl;
     }
     
     static void Error(const std::string& message) {
         std::cerr << "[ERROR] " << message << std::endl;
     }
 };
 
 // RAII wrapper for ACL resources
 class AclContext {
 public:
     AclContext(int32_t deviceId = DEFAULT_DEVICE_ID) : deviceId_(deviceId) {}
     
     ~AclContext() {
         Release();
     }
     
     GeluError Initialize() {
         aclError ret = aclInit(nullptr);
         if (ret != ACL_SUCCESS) {
             Logger::Error("aclInit failed: " + std::to_string(ret));
             return GeluError::ACL_INIT_FAILURE;
         }
         
         ret = aclrtSetDevice(deviceId_);
         if (ret != ACL_SUCCESS) {
             Logger::Error("aclrtSetDevice failed: " + std::to_string(ret));
             return GeluError::DEVICE_SET_FAILURE;
         }
         
         ret = aclrtCreateStream(&stream_);
         if (ret != ACL_SUCCESS) {
             Logger::Error("aclrtCreateStream failed: " + std::to_string(ret));
             return GeluError::STREAM_CREATE_FAILURE;
         }
         
         return GeluError::SUCCESS;
     }
     
     void Release() {
         if (stream_) {
             aclrtDestroyStream(stream_);
             stream_ = nullptr;
         }
         aclrtResetDevice(deviceId_);
         aclFinalize();
     }
     
     aclrtStream GetStream() const { return stream_; }
     
 private:
     int32_t deviceId_;
     aclrtStream stream_ = nullptr;
 };
 
 // File operations
 class FileUtil {
 public:
     static GeluError ReadBinaryFile(const std::string& filePath, std::vector<uint8_t>& buffer) {
         std::ifstream file(filePath, std::ios::binary | std::ios::ate);
         if (!file) {
             Logger::Error("Failed to open file: " + filePath);
             return GeluError::FILE_OPERATION_FAILURE;
         }
         
         std::streamsize size = file.tellg();
         file.seekg(0, std::ios::beg);
         
         buffer.resize(size);
         if (!file.read(reinterpret_cast<char*>(buffer.data()), size)) {
             Logger::Error("Failed to read file: " + filePath);
             return GeluError::FILE_OPERATION_FAILURE;
         }
         
         return GeluError::SUCCESS;
     }
     
     static GeluError WriteBinaryFile(const std::string& filePath, const void* data, size_t size) {
         std::ofstream file(filePath, std::ios::binary);
         if (!file) {
             Logger::Error("Failed to create file: " + filePath);
             return GeluError::FILE_OPERATION_FAILURE;
         }
         
         if (!file.write(reinterpret_cast<const char*>(data), size)) {
             Logger::Error("Failed to write file: " + filePath);
             return GeluError::FILE_OPERATION_FAILURE;
         }
         
         return GeluError::SUCCESS;
     }
 };
 
 // Tensor wrapper
 class AclTensor {
 public:
     AclTensor(const std::vector<int64_t>& shape, aclDataType dtype, void* data) 
         : shape_(shape), dtype_(dtype), data_(data) {
         tensor_ = aclCreateTensor(shape.data(), shape.size(), dtype, 
                                 nullptr, 0, ACL_FORMAT_ND,
                                 shape.data(), shape.size(), data);
         if (!tensor_) {
             Logger::Error("Failed to create ACL tensor");
         }
     }
     
     ~AclTensor() {
         if (tensor_) {
             aclDestroyTensor(tensor_);
         }
     }
     
     aclTensor* Get() const { return tensor_; }
     size_t GetSizeInBytes() const {
         size_t size = 1;
         for (auto dim : shape_) {
             size *= dim;
         }
         return size * aclDataTypeSize(dtype_);
     }
     
 private:
     std::vector<int64_t> shape_;
     aclDataType dtype_;
     void* data_;
     aclTensor* tensor_;
 };
 
 // Memory management
 class DeviceMemory {
 public:
     DeviceMemory(size_t size) {
         aclError ret = aclrtMalloc(&ptr_, size, ACL_MEM_MALLOC_HUGE_FIRST);
         if (ret != ACL_SUCCESS) {
             Logger::Error("Failed to allocate device memory: " + std::to_string(ret));
             ptr_ = nullptr;
         }
     }
     
     ~DeviceMemory() {
         if (ptr_) {
             aclrtFree(ptr_);
         }
     }
     
     void* Get() const { return ptr_; }
     operator bool() const { return ptr_ != nullptr; }
     
 private:
     void* ptr_ = nullptr;
 };
 
 } // namespace
 
 int main() {
     // 1. Initialize ACL context
     AclContext context;
     if (auto err = context.Initialize(); err != GeluError::SUCCESS) {
         return static_cast<int>(err);
     }
     
     // 2. Define tensor shapes
     const std::vector<int64_t> inputShape = {1, 100};
     const std::vector<int64_t> outputShape = {1, 100};
     
     // 3. Load input data
     std::vector<uint8_t> inputData;
     if (auto err = FileUtil::ReadBinaryFile("../input/input_x.bin", inputData); 
         err != GeluError::SUCCESS) {
         return static_cast<int>(err);
     }
     
     // 4. Allocate device memory
     DeviceMemory inputMem(inputData.size());
     DeviceMemory outputMem(outputShape[0] * outputShape[1] * FLOAT16_SIZE);
     if (!inputMem || !outputMem) {
         return static_cast<int>(GeluError::MEMORY_ALLOCATION_FAILURE);
     }
     
     // 5. Copy input data to device
     aclError ret = aclrtMemcpy(inputMem.Get(), inputData.size(), 
                               inputData.data(), inputData.size(),
                               ACL_MEMCPY_HOST_TO_DEVICE);
     if (ret != ACL_SUCCESS) {
         Logger::Error("Failed to copy input data to device: " + std::to_string(ret));
         return static_cast<int>(GeluError::MEMORY_COPY_FAILURE);
     }
     
     // 6. Create tensors
     AclTensor inputTensor(inputShape, DATA_TYPE, inputMem.Get());
     AclTensor outputTensor(outputShape, DATA_TYPE, outputMem.Get());
     if (!inputTensor.Get() || !outputTensor.Get()) {
         return static_cast<int>(GeluError::TENSOR_CREATION_FAILURE);
     }
     
     // 7. Execute GELU operator
     uint64_t workspaceSize = 0;
     aclOpExecutor* executor = nullptr;
     
     ret = aclnnGeluGetWorkspaceSize(inputTensor.Get(), outputTensor.Get(), 
                                    &workspaceSize, &executor);
     if (ret != ACL_SUCCESS) {
         Logger::Error("Failed to get workspace size: " + std::to_string(ret));
         return static_cast<int>(GeluError::WORKSPACE_QUERY_FAILURE);
     }
     
     std::unique_ptr<DeviceMemory> workspace;
     if (workspaceSize > 0) {
         workspace = std::make_unique<DeviceMemory>(workspaceSize);
         if (!*workspace) {
             return static_cast<int>(GeluError::MEMORY_ALLOCATION_FAILURE);
         }
     }
     
     ret = aclnnGelu(workspace ? workspace->Get() : nullptr, 
                    workspaceSize, executor, context.GetStream());
     if (ret != ACL_SUCCESS) {
         Logger::Error("Failed to execute GELU operator: " + std::to_string(ret));
         return static_cast<int>(GeluError::OPERATOR_EXECUTION_FAILURE);
     }
     
     // 8. Synchronize and get results
     ret = aclrtSynchronizeStream(context.GetStream());
     if (ret != ACL_SUCCESS) {
         Logger::Error("Failed to synchronize stream: " + std::to_string(ret));
         return static_cast<int>(GeluError::STREAM_SYNC_FAILURE);
     }
     
     std::vector<aclFloat16> outputData(outputTensor.GetSizeInBytes() / FLOAT16_SIZE);
     ret = aclrtMemcpy(outputData.data(), outputData.size() * FLOAT16_SIZE,
                      outputMem.Get(), outputTensor.GetSizeInBytes(),
                      ACL_MEMCPY_DEVICE_TO_HOST);
     if (ret != ACL_SUCCESS) {
         Logger::Error("Failed to copy output data: " + std::to_string(ret));
         return static_cast<int>(GeluError::MEMORY_COPY_FAILURE);
     }
     
     // 9. Save output
     if (auto err = FileUtil::WriteBinaryFile("../output/output_z.bin", 
                                            outputData.data(), 
                                            outputData.size() * FLOAT16_SIZE);
         err != GeluError::SUCCESS) {
         return static_cast<int>(err);
     }
     
     Logger::Info("GELU operation completed successfully");
     return static_cast<int>(GeluError::SUCCESS);
 }
    